from judge import command, commands
from config import TESTDATA_PATH, MAX_JUDGE_THREAD
from runner import Testcase, compile_exec, Result
from concurrent.futures import ThreadPoolExecutor
import tempfile
import os
import shutil
import zipfile
import runner
import threading


@command(name="help", help="显示本页面", usage="help ")
def help(args):
    text = "BZOJ-Judger\n已注册命令:\n"
    for name, value in commands.items():
        text += "{}   ---   {}   ---   {}\n".format(name, value[0], value[1])
    print(text)


@command(name="single", help="进行单点评测", usage="single [程序文件名] [题目编号] [输入文件] [输出文件]")
def single_test(args: list):
    source, problem_id, input_file, output_file = args
    judge_dir = tempfile.mkdtemp()
    testdata_zip: zipfile.ZipFile = get_problem_zip(problem_id)

    def path(p):
        return os.path.join(judge_dir, p)
    for file in testdata_zip.namelist():
        file_name = file.split("/")[-1]
        if file_name == input_file or file_name == output_file:
            with open(path(file_name), "wb") as output:
                with testdata_zip.open(file) as input:
                    output.write(input.read())
    shutil.copy(source, path(source))
    compile_exec(judge_dir, source,
                 source+".exe")
    result: Result = runner.run_standard(source+".exe", input_file, output_file,
                                         judge_dir, 998244353, 998244353)
    print("\n输入:{}\n你的输出:{}\n答案:{}\n{}".format("".join(result[1][0]), "".join(
        result[1][1]), "".join(result[1][2]), result[0]))


@command(name="copy", help="复制测试文件到当前目录", usage="copt [题目编号] [输入文件名] [输出文件名]")
def copy_data(args):
    problem_id, input_file, output_file = args
    testdata_zip: zipfile.ZipFile = get_problem_zip(problem_id)
    def path(p):
        return os.path.join(judge_dir, p)
    for file in testdata_zip.namelist():
        file_name = file.split("/")[-1]
        if file_name == input_file or file_name == output_file:
            with open(file_name, "wb") as output:
                with testdata_zip.open(file) as input:
                    output.write(input.read())


@command(name="judge", help="进行评测", usage="judge [程序文件名] [题目编号] [时限(默认1000ms)] [内存限制(默认512M)]")
def judge(args):
    source, problem_id = args[:2]
    time_limit = 1000
    memory_limit = 512
    if len(args) >= 3:
        time_limit = int(args[2])
    if len(args) >= 4:
        memory_limit = int(args[3])
    print("时限:{}ms\n内存限制:{}MB".format(time_limit, memory_limit))
    problem_id = int(problem_id)
    judge_dir = tempfile.mkdtemp()
    source_target = os.path.join(
        judge_dir, source if "/" not in source else source.split("/")[-1])
    print("复制源代码到:{}".format(source_target))
    shutil.copyfile(source, source_target)
    print("解压测试数据中")
    testdata_zip: zipfile.ZipFile = get_problem_zip(problem_id)
    print(testdata_zip.namelist())
    testdata_list: list = unzip_testdata(judge_dir, testdata_zip)
    compile_result = compile_exec(judge_dir, source,
                                  source+".exe")
    print("编译结果:\n{}".format(compile_result))
    results = []

    def run(pack):
        index, testdata = pack
        temp = runner.run_standard(source+".exe", testdata.input, testdata.output,
                                   judge_dir, time_limit, memory_limit)
        print("测试点 {}/{} 已提交".format(index+1, len(testdata_list)))
        if temp is TypeError or temp is Exception:
            raise temp
        else:
            return temp[0]
    # QQ:489679970
    poorpool = ThreadPoolExecutor(MAX_JUDGE_THREAD, "JUDGER")
    results = list(poorpool.map(run, enumerate(testdata_list)))
    poorpool.shutdown()
    sorted(results, key=lambda x: x[0])
    total = 0
    status = {}
    for index, data in enumerate(results):
        # print("index={} len={}".format(index,len(testdata_list)))
        print("测试点{}({}/{}):{}".format(index+1,
                                       testdata_list[index].input, testdata_list[index].output, data))
        total += data.score
        curr = data.status
        if curr in status:
            status[curr] += 1
        else:
            status[curr] = 1
    print("总分:{}".format(total/(len(testdata_list))))
    print("共计{}个测试点，其中:".format(len(testdata_list)))
    for key, value in status.items():
        print("{} : {} ({}%)".format(key, value, value/len(testdata_list)*100))


def get_problem_zip(problem_id)->zipfile.ZipFile:
    return zipfile.ZipFile(os.path.join(
        TESTDATA_PATH, "{}.zip".format(problem_id)))


def unzip_testdata(judge_dir, testdata_zip)->list:
    testdata_list = []
    ins = []
    outs = []
    for file in filter(lambda x: not x.endswith("/"), testdata_zip.namelist()):
        output = open(os.path.join(
            judge_dir, file[file.index("/")+1:]), "wb")
        input = testdata_zip.open(file)
        output.write(input.read())
        input.close()
        output.close()
        if file.endswith("in"):
            ins.append(file[file.index("/")+1:])
        else:
            outs.append(file[file.index("/")+1:])
    ins = sorted(ins)
    outs = sorted(outs)
    for i in range(0, len(ins)):
        testdata_list.append(Testcase(ins[i], outs[i]))
    testdata_zip.close()
    return testdata_list
